#include "../../include/Instruction.h"
#include "../../include/ir.h"

#include <assert.h>
#include <set>
#include <stack>
extern std::map<FuncDefInstruction, int> max_label_map;
extern std::map<FuncDefInstruction, int> max_reg_map;

std::map<std::string, CFG *> CFGMap;

void LLVMIR::CFGInit() {
    for (auto &[defI, bb_map] : function_block_map) {
        CFG *cfg = new CFG();
        cfg->block_map = &bb_map;
        cfg->function_def = defI;                
        cfg->max_reg = max_reg_map[defI];      
        cfg->max_label = max_label_map[defI];
        cfg->BuildCFG();
        //TODO("init your members in class CFG if you need");
        CFGMap[defI->GetFunctionName()] = cfg;
        llvm_cfg[defI] = cfg;
    }
}

void LLVMIR::BuildCFG() {
    for (auto [defI, cfg] : llvm_cfg) {
        cfg->BuildCFG();
    }
}

void EliminateUnreachedBlocksInsts(CFG *C) {
    std::queue<int> q;
    q.push(0);
    std::vector<bool> visited(C->block_map->size(), false);
    while(!q.empty()){
        int bbid = q.front();
        visited[bbid] = true;      
        q.pop();
        BasicBlock *bb = C->block_map->at(bbid);
        for (auto it = bb->Instruction_list.begin(); it != bb->Instruction_list.end(); ++it) {
            BasicInstruction *ins = *it;
            if (ins->GetOpcode() ==BasicInstruction:: BR_UNCOND ){
                BrUncondInstruction *bruncond = (BrUncondInstruction *)ins;
                int target_block_no = ((LabelOperand *)bruncond->GetDestLabel())->GetLabelNo();
                if(visited[target_block_no] == false) {
                q.push(target_block_no);
                }  
                bb->Instruction_list.erase(it + 1, bb->Instruction_list.end());
                break;
            }
            else if (ins->GetOpcode() ==BasicInstruction:: BR_COND){
                BrCondInstruction *brcond = (BrCondInstruction *)ins;
                int target_trueblock_no = ((LabelOperand *)brcond->GetTrueLabel())->GetLabelNo();
                int target_falseblock_no = ((LabelOperand *)brcond->GetFalseLabel())->GetLabelNo();
                if(visited[target_trueblock_no] == false) {
                q.push(target_trueblock_no);
                } 
                if(visited[target_falseblock_no] == false) {
                q.push(target_falseblock_no);
                }

                bb->Instruction_list.erase(it + 1, bb->Instruction_list.end());
                break;
            }
            else if( ins->GetOpcode() == BasicInstruction::RET) {
                // 如果是跳转或返回指令，删除该指令后的其余指令
                bb->Instruction_list.erase(it + 1, bb->Instruction_list.end());
                break;
            }
        }
    }
    int size = C->block_map->size();
    for (int it=0;it<size;it++) {
        //std::cout << "it->first: " << it->first << std::endl;
        if (!visited[it]) {
            //std::cout << "erase bbid: " << it << std::endl;
            C->block_map->erase(it); // 删除基本块            
        } 
    }
}
//reference https://github.com/optimize\cfg\cfg.cc 103-129
void CFG::BuildCFG() { 
    EliminateUnreachedBlocksInsts(this);
    G.clear();
    invG.clear();
    G.resize(max_label + 1);
    invG.resize(max_label + 1);
    for (auto [id, bb] : *block_map) {
        Instruction lastIns = bb->Instruction_list[bb->Instruction_list.size() - 1];
        if (lastIns->GetOpcode() ==BasicInstruction::BR_UNCOND) {
            BrUncondInstruction *bruncond = (BrUncondInstruction *)lastIns;
            int target_block_no = ((LabelOperand *)bruncond->GetDestLabel())->GetLabelNo();
            G[id].push_back((*block_map)[target_block_no]);
            invG[target_block_no].push_back(bb);
        } else if (lastIns->GetOpcode() == BasicInstruction::BR_COND) {
            BrCondInstruction *brcond = (BrCondInstruction *)lastIns;
            int target_trueblock_no = ((LabelOperand *)brcond->GetTrueLabel())->GetLabelNo();
            int target_falseblock_no = ((LabelOperand *)brcond->GetFalseLabel())->GetLabelNo();
            G[id].push_back((*block_map)[target_trueblock_no]);     
            G[id].push_back((*block_map)[target_falseblock_no]);    
                  
            invG[target_trueblock_no].push_back(bb);
            
            invG[target_falseblock_no].push_back(bb);
           
        } else if (lastIns->GetOpcode() == BasicInstruction::RET) {
            ret_block = bb;
        }
    }
}

std::vector<LLVMBlock> CFG::GetPredecessor(LLVMBlock B) { return invG[B->block_id]; }

std::vector<LLVMBlock> CFG::GetPredecessor(int bbid) { return invG[bbid]; }

std::vector<LLVMBlock> CFG::GetSuccessor(LLVMBlock B) { return G[B->block_id]; }

std::vector<LLVMBlock> CFG::GetSuccessor(int bbid) { return G[bbid]; }